
#include <string.h>
#include <stdlib.h>
#include "vfs.h"
#include "vfs_lfs.h"
#include "sysdef.h"
#include "device.h"

#define LFS_BLOCK_CYCLES    500
#define LFS_LOOKAHEAD_SIZE  16

#define LFS_OPEN_TXT_FLAG   0x1000

extern vfs_t __lfsRegTable_start[];
extern vfs_t __lfsRegTable_end[];
extern char __flash_user_start[];
extern char __flash_user_end[];

static int lfs_flash_read(const struct lfs_config* c, lfs_block_t block, lfs_off_t off, void* buffer, lfs_size_t size) {
    flash_type_e flash_type = FLASH_TYPE_MAX;

#if CONFIG_LFS_USE_FLASH_ONCHIP
    if(((lfs_cfg_ctx_t *)(c->context))->start_addr >= (int)__flash_user_start && ((lfs_cfg_ctx_t *)(c->context))->end_addr <= (int)__flash_user_end) {
        flash_type = FLASH_TYPE_ONCHIP;
    }
#endif

#if CONFIG_LFS_USE_FLASH_W25QXX
    if(((lfs_cfg_ctx_t *)(c->context))->start_addr < (int)__flash_user_start || ((lfs_cfg_ctx_t *)(c->context))->end_addr > (int)__flash_user_end) {
        flash_type = ((lfs_cfg_ctx_t *)(c->context))->start_addr / 0x01000000 + FLASH_TYPE_W25QXX;
    }
#endif

    flash_t *flash = flash_find_by_type(flash_type);
    uint32_t lfs_start_address = ((lfs_cfg_ctx_t *)(c->context))->start_addr;
    uint32_t read_addr = block * c->block_size + lfs_start_address + off;
    return flash->read(read_addr, (uint8_t *)buffer, (uint16_t)size);
}

static int lfs_flash_prog(const struct lfs_config* c, lfs_block_t block, lfs_off_t off, const void* buffer, lfs_size_t size) {
    flash_type_e flash_type = FLASH_TYPE_MAX;

#if CONFIG_LFS_USE_FLASH_ONCHIP
    if(((lfs_cfg_ctx_t *)(c->context))->start_addr >= (int)__flash_user_start && ((lfs_cfg_ctx_t *)(c->context))->end_addr <= (int)__flash_user_end) {
        flash_type = FLASH_TYPE_ONCHIP;
    }
#endif

#if CONFIG_LFS_USE_FLASH_W25QXX
    if(((lfs_cfg_ctx_t *)(c->context))->start_addr < (int)__flash_user_start || ((lfs_cfg_ctx_t *)(c->context))->end_addr > (int)__flash_user_end) {
        flash_type = ((lfs_cfg_ctx_t *)(c->context))->start_addr / 0x01000000 + FLASH_TYPE_W25QXX;
    }
#endif

    flash_t *flash = flash_find_by_type(flash_type);
    uint32_t lfs_start_address = ((lfs_cfg_ctx_t *)(c->context))->start_addr;
    uint32_t write_addr = block * c->block_size + lfs_start_address + off;
    return flash->write(write_addr, (uint8_t *)buffer, (uint16_t)size);
}

static int lfs_flash_erase(const struct lfs_config* c, lfs_block_t block) {
    flash_type_e flash_type = FLASH_TYPE_MAX;

#if CONFIG_LFS_USE_FLASH_ONCHIP
    if(((lfs_cfg_ctx_t *)(c->context))->start_addr >= (int)__flash_user_start && ((lfs_cfg_ctx_t *)(c->context))->end_addr <= (int)__flash_user_end) {
        flash_type = FLASH_TYPE_ONCHIP;
    }
#endif

#if CONFIG_LFS_USE_FLASH_W25QXX
    if(((lfs_cfg_ctx_t *)(c->context))->start_addr < (int)__flash_user_start || ((lfs_cfg_ctx_t *)(c->context))->end_addr > (int)__flash_user_end) {
        flash_type = ((lfs_cfg_ctx_t *)(c->context))->start_addr / 0x01000000 + FLASH_TYPE_W25QXX;
    }
#endif

    flash_t *flash = flash_find_by_type(flash_type);
    uint32_t lfs_start_address = ((lfs_cfg_ctx_t *)(c->context))->start_addr;
    uint32_t erase_addr = block * c->block_size + lfs_start_address;
    return flash->erase(erase_addr, block * c->block_size);
}

static int lfs_flash_sync(const struct lfs_config* c) {
    UNUSED(c);
    return LFS_ERR_OK;
}

static int lfs_lock(const struct lfs_config* c) {
    OS_ERR err;
    OSMutexPend(&((lfs_cfg_ctx_t *)(c->context))->mutex, OS_WAIT_FOREVER, OS_OPT_PEND_NON_BLOCKING, NULL, &err);
    return 0;
}

static int lfs_unlock(const struct lfs_config* c) {
    OS_ERR err;
    OSMutexPost(&((lfs_cfg_ctx_t *)(c->context))->mutex, OS_OPT_POST_NONE, &err);
    return 0;
}

static const struct {
    char *mode;
    int flags;
} flags_map[] = {
    {"r",   LFS_OPEN_FLAGS_R},
    {"rb",  LFS_OPEN_FLAGS_R},
    {"rt",  LFS_OPEN_FLAGS_R},
    {"r+",  LFS_OPEN_FLAGS_RPluse},
    {"r+b", LFS_OPEN_FLAGS_RPluse},
    {"r+t", LFS_OPEN_FLAGS_RPluse},
    {"rb+", LFS_OPEN_FLAGS_RPluse},
    {"rt+", LFS_OPEN_FLAGS_RPluse},
    {"w",   LFS_OPEN_FLAGS_W},
    {"wb",  LFS_OPEN_FLAGS_W},
    {"wt",  LFS_OPEN_FLAGS_W},
    {"w+",  LFS_OPEN_FLAGS_WPluse},
    {"w+b", LFS_OPEN_FLAGS_WPluse},
    {"w+t", LFS_OPEN_FLAGS_WPluse},
    {"wb+", LFS_OPEN_FLAGS_WPluse},
    {"wt+", LFS_OPEN_FLAGS_WPluse},
    {"a",   LFS_OPEN_FLAGS_A},
    {"ab",  LFS_OPEN_FLAGS_A},
    {"at",  LFS_OPEN_FLAGS_A},
    {"a+",  LFS_OPEN_FLAGS_APluse},
    {"a+b", LFS_OPEN_FLAGS_APluse},
    {"a+t", LFS_OPEN_FLAGS_APluse},
    {"ab+", LFS_OPEN_FLAGS_APluse},
    {"at+", LFS_OPEN_FLAGS_APluse},
    {NULL, 0}
};

static int lfs_mode2flags(const char *mode) {
    int i = 0;
    while(flags_map[i].mode) {
        if(strcasecmp(flags_map[i].mode, mode) == 0) return flags_map[i].flags;
        i++;
    }

    return 0;
}

static vfs_t *lfs_get_by_path(const char *path, char **lfs_path, char **abs_path) {
    if(!path || !lfs_path || !abs_path) return NULL;

    vfs_t *vfs = vfs_get_by_path(path, lfs_path, abs_path);
    if(!vfs) return NULL;

    if(vfs->type != VFS_TYPE_LFS) {
        free(*abs_path);
        *abs_path = NULL;
        return NULL;
    }

    return vfs;
}

static FILE *vfs_lfs_fopen(const char *path, const char *mode) {
    vfs_t *vfs = NULL;
    lfs_t *lfs = NULL;
    FILE *fp = NULL;
    lfs_file_t *file = NULL;
    int flags = 0;
    char *lfs_path = NULL;
    char *abs_path = NULL;

    if(!path || !mode) goto err;

    vfs = lfs_get_by_path(path, &lfs_path, &abs_path);
    if(!vfs) goto err;

    lfs = &vfs->u.lfs_obj.lfs;

    flags = lfs_mode2flags(mode);
    if(!flags) goto err;

    fp = malloc(sizeof(FILE));
    if(!fp) goto err;
    memset(fp, 0, sizeof(FILE));
    
    file = malloc(sizeof(lfs_file_t));
    if(!file) goto err;
    memset(file, 0, sizeof(lfs_file_t));

    if(lfs_file_open(lfs, file, lfs_path, flags))
        goto err;

    if(!strchr(mode, 'b'))
        flags |= LFS_OPEN_TXT_FLAG;
    fp->_flags = (short)flags;
    fp->_flags2 = (int)file;
    fp->_cookie = (void *)vfs;
    if(abs_path) free(abs_path);
    return fp;

err:
    if(fp) free(fp);
    if(file) free(file);
    if(abs_path) free(abs_path);
    return NULL;
}

static int vfs_lfs_fclose(FILE * fp) {
    if(!fp) goto err;

    vfs_t *vfs = (vfs_t *)fp->_cookie;
    lfs_t *lfs = &vfs->u.lfs_obj.lfs;
    lfs_file_t *file = (lfs_file_t *)fp->_flags2;

    if(lfs_file_close(lfs, file)) goto err;

    free(file);
    free(fp);
    return 0;

err:
    return -1;
}

static size_t vfs_lfs_fread(void *__restrict buf, size_t size, size_t count, FILE *__restrict fp) {
    if(!fp) goto err;

    vfs_t *vfs = (vfs_t *)fp->_cookie;
    lfs_t *lfs = &vfs->u.lfs_obj.lfs;
    lfs_file_t *file = (lfs_file_t *)fp->_flags2;

    lfs_ssize_t nread = lfs_file_read(lfs, file, buf, (lfs_size_t)(size*count));
    if(nread < 0) goto err;

    if(fp->_flags & LFS_OPEN_TXT_FLAG) {
        lfs_ssize_t n = 0;
        while(((char *)buf)[n] && n < nread) n++;
        nread = n;
    }

    return nread;

err:
    return -1;
}

static size_t vfs_lfs_fwrite(const void *__restrict buf, size_t size, size_t count, FILE * fp) {
    if(!fp) goto err;

    vfs_t *vfs = (vfs_t *)fp->_cookie;
    lfs_t *lfs = &vfs->u.lfs_obj.lfs;
    lfs_file_t *file = (lfs_file_t *)fp->_flags2;

    lfs_ssize_t nwrite = size*count;
    if(fp->_flags & LFS_OPEN_TXT_FLAG) {
        lfs_ssize_t n = 0;
        while(((char *)buf)[n] && n < nwrite) n++;
        nwrite = n;
    }

    nwrite = lfs_file_write(lfs, file, buf, (lfs_size_t)nwrite);
    if(nwrite < 0) goto err;

    return nwrite;

err:
    return -1;
}

static int vfs_lfs_fflush(FILE *fp) {
    if(!fp) return -1;

    vfs_t *vfs = (vfs_t *)fp->_cookie;
    lfs_t *lfs = &vfs->u.lfs_obj.lfs;
    lfs_file_t *file = (lfs_file_t *)fp->_flags2;

    if(lfs_file_sync(lfs, file)) return -1;

    return 0;
}

static int vfs_lfs_fseek(FILE *fp, long offset, int whence) {
    if(!fp) return -1;

    vfs_t *vfs = (vfs_t *)fp->_cookie;
    lfs_t *lfs = &vfs->u.lfs_obj.lfs;
    lfs_file_t *file = (lfs_file_t *)fp->_flags2;

    return lfs_file_seek(lfs, file, offset, whence) ? -1 : 0;
}

static long vfs_lfs_ftell(FILE * fp) {
    if(!fp) return -1;

    vfs_t *vfs = (vfs_t *)fp->_cookie;
    lfs_t *lfs = &vfs->u.lfs_obj.lfs;
    lfs_file_t *file = (lfs_file_t *)fp->_flags2;

    long pos = lfs_file_tell(lfs, file);
    if(pos < 0) pos = -1;

    return pos;
}

static void vfs_lfs_rewind(FILE * fp) {
    if(fp) {
        vfs_t *vfs = (vfs_t *)fp->_cookie;
        lfs_t *lfs = &vfs->u.lfs_obj.lfs;
        lfs_file_t *file = (lfs_file_t *)fp->_flags2;
        lfs_file_rewind(lfs, file);
    }
}

static int vfs_lfs_mkdir(const char *path, mode_t mode) {
    UNUSED(mode);
    
    if(!path) goto err;

    char *lfs_path = NULL;
    char *abs_path = NULL;

    vfs_t *vfs = lfs_get_by_path(path, &lfs_path, &abs_path);
    if(!vfs) goto err;

    lfs_t *lfs = &vfs->u.lfs_obj.lfs;

    if(lfs_mkdir(lfs, lfs_path)) goto err;

    if(abs_path) free(abs_path);
    return 0;

err:
    if(abs_path) free(abs_path);
    return -1;
}

static DIR *vfs_lfs_opendir(const char *path) {
    vfs_t *vfs = NULL;
    lfs_t *lfs = NULL;
    DIR *dp = NULL;
    lfs_dir_t *dir = NULL;
    char *lfs_path = NULL;
    char *abs_path = NULL;

    if(!path) goto err;

    vfs = lfs_get_by_path(path, &lfs_path, &abs_path);
    if(!vfs) goto err;

    lfs = &vfs->u.lfs_obj.lfs;

    dp = malloc(sizeof(DIR));
    if(!dp) goto err;
    memset(dp, 0, sizeof(DIR));

    dir = malloc(sizeof(lfs_dir_t));
    if(!dir) goto err;

    if(lfs_dir_open(lfs, dir, lfs_path)) goto err;

    dp->dd_fd = (int)dir;
    dp->dd_loc = (int)vfs;
    if(abs_path) free(abs_path);
    return dp;

err:
    if(dp) free(dp);
    if(dir) free(dir);
    if(abs_path) free(abs_path);
    return NULL;
}

static int vfs_lfs_closedir(DIR *dp) {
    if(!dp) goto err;

    vfs_t *vfs = (vfs_t *)dp->dd_loc;
    lfs_t *lfs = &vfs->u.lfs_obj.lfs;
    lfs_dir_t *dir = (lfs_dir_t *)dp->dd_fd;

    if(lfs_dir_close(lfs, dir)) goto err;

    free(dir);
    free(dp);
    return 0;

err:
    return -1;
}

static struct dirent *vfs_lfs_readdir(DIR *dp) {
    if(!dp) goto err;

    vfs_t *vfs = (vfs_t *)dp->dd_loc;
    lfs_t *lfs = &vfs->u.lfs_obj.lfs;
    lfs_dir_t *dir = (lfs_dir_t *)dp->dd_fd;
    struct lfs_info info = {0};

    if(lfs_dir_read(lfs, dir, &info) <= 0) goto err;

    static struct dirent dirent = {0};
    memset(&dirent, 0, sizeof(struct dirent));
    dirent.d_fileno = info.size;
    dirent.d_reclen = info.type;
    dirent.d_namlen = strlen(info.name);
    strcpy(dirent.d_name, info.name);
    return &dirent;

err:
    return NULL;
}

static void vfs_lfs_seekdir(DIR *dp, long loc) {
    if(!dp) return;

    vfs_t *vfs = (vfs_t *)dp->dd_loc;
    lfs_t *lfs = &vfs->u.lfs_obj.lfs;
    lfs_dir_t *dir = (lfs_dir_t *)dp->dd_fd;

    lfs_dir_seek(lfs, dir, loc);
}

static long vfs_lfs_telldir(DIR *dp) {
    if(!dp) return -1;

    vfs_t *vfs = (vfs_t *)dp->dd_loc;
    lfs_t *lfs = &vfs->u.lfs_obj.lfs;
    lfs_dir_t *dir = (lfs_dir_t *)dp->dd_fd;

    long pos = lfs_dir_tell(lfs, dir);
    if(pos < 0) pos = -1;

    return pos;
}

static void vfs_lfs_rewinddir(DIR *dp) {
    if(!dp) return;

    vfs_t *vfs = (vfs_t *)dp->dd_loc;
    lfs_t *lfs = &vfs->u.lfs_obj.lfs;
    lfs_dir_t *dir = (lfs_dir_t *)dp->dd_fd;

    lfs_dir_rewind(lfs, dir);
}

static int vfs_lfs_rmdir(const char *path) {
    char *lfs_path = NULL;
    char *abs_path = NULL;

    if(!path) goto err;

    vfs_t *vfs = lfs_get_by_path(path, &lfs_path, &abs_path);
    if(!vfs) goto err;

    lfs_t *lfs = &vfs->u.lfs_obj.lfs;

    struct lfs_info info = {0};
    if(lfs_stat(lfs, lfs_path, &info)) goto err;
    if(info.type != LFS_TYPE_DIR) goto err;

    if(lfs_remove(lfs, lfs_path)) goto err;

    if(abs_path) free(abs_path);
    return 0;

err:
    if(abs_path) free(abs_path);
    return -1;
}

static int vfs_lfs_stat(const char *path, struct stat *buf) {
    char *lfs_path = NULL;
    char *abs_path = NULL;

    if(!path || !buf) goto err;

    vfs_t *vfs = lfs_get_by_path(path, &lfs_path, &abs_path);
    if(!vfs) goto err;

    lfs_t *lfs = &vfs->u.lfs_obj.lfs;

    struct lfs_info info = {0};
    if(lfs_stat(lfs, lfs_path, &info)) goto err;

    memset(buf, 0, sizeof(struct stat));

    if(info.type == LFS_TYPE_REG) {
        buf->st_mode = S_IFREG;

        FILE *fp = vfs_lfs_fopen(path, "rb");
        if(!fp) goto err;

        lfs_file_t *file = (lfs_file_t *)fp->_flags2;
        buf->st_size = lfs_file_size(lfs, file);

        if(buf->st_size < 0) {
            buf->st_size = 0;
            vfs_lfs_fclose(fp);
            goto err;
        }

        buf->st_blksize = lfs->cfg->block_size;
        buf->st_blocks = file->block;
        buf->st_ino = file->off;
        vfs_lfs_fclose(fp);
    }
    else {
        buf->st_mode = S_IFDIR;

        DIR *dp = vfs_lfs_opendir(path);
        if(!dp) goto err;

        lfs_dir_t *dir = (lfs_dir_t *)dp->dd_fd;
        buf->st_ino = dir->pos;
        vfs_lfs_closedir(dp);
    }

    if(abs_path) free(abs_path);
    return 0;

err:
    if(abs_path) free(abs_path);
    return -1;
}

static int vfs_lfs_access(const char *path, int mode) {
    UNUSED(mode);

    if(!path) return -1;

    struct stat buf;
    if(vfs_lfs_stat(path, &buf))
        return -1;
    else
        return 0;
}

static int vfs_lfs_remove(const char *path) {
    char *lfs_path = NULL;
    char *abs_path = NULL;

    if(!path) goto err;

    vfs_t *vfs = lfs_get_by_path(path, &lfs_path, &abs_path);
    if(!vfs) goto err;

    lfs_t *lfs = &vfs->u.lfs_obj.lfs;

    if(lfs_remove(lfs, lfs_path)) goto err;

    if(abs_path) free(abs_path);
    return 0;

err:
    if(abs_path) free(abs_path);
    return -1;
}

static int vfs_lfs_rename(const char *oldpath, const char *newpath) {
    if(!oldpath || !newpath) goto err;

    char * old_lfs_path = NULL, * new_lfs_path = NULL;
    char * old_abs_path = NULL, * new_abs_path = NULL;
    FILE * old_fp = NULL, * new_fp = NULL;
    uint8_t *buf = NULL;
    int ret = 0;

    vfs_t * old_vfs = lfs_get_by_path(oldpath, &old_lfs_path, &old_abs_path);
    if(!old_vfs) {
        ret = -1;
        goto end;
    }

    lfs_t *old_lfs = &old_vfs->u.lfs_obj.lfs;

    vfs_t * new_vfs = lfs_get_by_path(newpath, &new_lfs_path, &new_abs_path);
    if(!new_vfs) {
        ret = -1;
        goto end;
    }

    lfs_t *new_lfs = &new_vfs->u.lfs_obj.lfs;

    if(old_lfs == new_lfs) {
        if(lfs_rename(old_lfs, old_lfs_path, new_lfs_path)) {
            ret = -1;
            goto end;
        }
    } else {
        if(vfs_lfs_access(newpath, 0) == 0) {
            ret = -1;
            goto end;
        }

        old_fp = vfs_lfs_fopen(oldpath, "rb");
        if(!old_fp) {
            ret = -1;
            goto end;
        }

        new_fp = vfs_lfs_fopen(newpath, "wb");
        if(!new_fp) {
            ret = -1;
            goto end;
        }

        buf = malloc(1024);
        if(!buf) {
            ret = -1;
            goto end;
        }

        int nread = 0;
        while((nread = vfs_lfs_fread(buf, 1, 1024, old_fp)) > 0) {
            if(vfs_lfs_fwrite(buf, 1, nread, new_fp) != nread) {
                ret = -1;
                goto end;
            }
        }

        if(nread < 0) {
            ret = -1;
            goto end;
        }
    }

end:
    if(old_fp) vfs_lfs_fclose(old_fp);
    if(new_fp) vfs_lfs_fclose(new_fp);
    if(buf) free(buf);

    if(old_abs_path) free(old_abs_path);
    if(new_abs_path) free(new_abs_path);

    if(ret != 0) {
        if(vfs_lfs_access(newpath, 0) == 0) {
            vfs_lfs_remove(newpath);
        }
        goto err;
    }

    return 0;

err:
    return -1;
}

const vfs_ops_t vfs_lfs_ops = {
    .fopen = vfs_lfs_fopen,
    .fclose = vfs_lfs_fclose,
    .fread = vfs_lfs_fread,
    .fwrite = vfs_lfs_fwrite,
    .fflush = vfs_lfs_fflush,
    .fseek = vfs_lfs_fseek,
    .ftell = vfs_lfs_ftell,
    .rewind = vfs_lfs_rewind,
    .mkdir = vfs_lfs_mkdir,
    .opendir = vfs_lfs_opendir,
    .closedir = vfs_lfs_closedir,
    .readdir = vfs_lfs_readdir,
    .seekdir = vfs_lfs_seekdir,
    .telldir = vfs_lfs_telldir,
    .rewinddir = vfs_lfs_rewinddir,
    .chdir = NULL,
    .getcwd = NULL,
    .rmdir = vfs_lfs_rmdir,
    .access = vfs_lfs_access,
    .stat = vfs_lfs_stat,
    .remove = vfs_lfs_remove,
    .rename = vfs_lfs_rename
};

int vfs_lfs_init(lfs_object_t *object, char *name) {
    flash_type_e flash_type = FLASH_TYPE_MAX;
    lfs_size_t read_size, prog_size, block_size;

#if CONFIG_LFS_USE_FLASH_ONCHIP
    if(object->cfg_ctx.start_addr >= (int)__flash_user_start && object->cfg_ctx.end_addr <= (int)__flash_user_end) {
        flash_type = FLASH_TYPE_ONCHIP;
    }
#endif

#if CONFIG_LFS_USE_FLASH_W25QXX
    if(object->cfg_ctx.start_addr < (int)__flash_user_start || object->cfg_ctx.end_addr > (int)__flash_user_end) {
        flash_type = object->cfg_ctx.start_addr / 0x01000000 + FLASH_TYPE_W25QXX;
    }
#endif

    flash_t *flash = flash_find_by_type(flash_type);
    if(!flash) return -1;

    if(flash->open(NULL)) {
        return -1;
    }

    read_size = prog_size = block_size = flash->info.block_size;

    if(object->cfg_ctx.end_addr - object->cfg_ctx.start_addr >= 4 * block_size
    && object->cfg_ctx.start_addr % block_size == 0 && object->cfg_ctx.end_addr % block_size == 0) {
        OS_ERR err;
        OSMutexCreate(&object->cfg_ctx.mutex, name, &err);
        if(err != OS_ERR_NONE) return -1;

        object->cfg.context = (void*)&object->cfg_ctx;

        object->cfg.lock = lfs_lock;
        object->cfg.unlock = lfs_unlock;

        object->cfg.read = lfs_flash_read;
        object->cfg.prog = lfs_flash_prog;
        object->cfg.erase = lfs_flash_erase;
        object->cfg.sync = lfs_flash_sync;

        object->cfg.read_size = read_size;
        object->cfg.prog_size = prog_size;
        object->cfg.block_size = block_size;
        object->cfg.block_cycles = LFS_BLOCK_CYCLES;
        object->cfg.block_count = (object->cfg_ctx.end_addr - object->cfg_ctx.start_addr) / block_size;
        object->cfg.cache_size = 4 * block_size;
        object->cfg.lookahead_size = LFS_LOOKAHEAD_SIZE;

        if(lfs_mount(&object->lfs, &object->cfg)) {
            if(lfs_format(&object->lfs, &object->cfg)) {
                return -1;
            }
            if(lfs_mount(&object->lfs, &object->cfg)) {
                return -1;
            }
        }
    }

    return 0;
}
